"""
文本分类算法
"""
import re
import jieba
import numpy as np
from typing import List, Dict, Tuple
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
from utils.logger import audit_logger
from models.schemas import TextClassificationResult


class TextClassifier:
    """文本分类器"""
    
    def __init__(self):
        self.model = None
        self.categories = {
            'normal': '正常内容',
            'spam': '垃圾信息',
            'advertisement': '广告推广',
            'inappropriate': '不当内容',
            'violence': '暴力内容',
            'adult': '成人内容'
        }
        self.init_model()
    
    def init_model(self):
        """初始化分类模型"""
        try:
            # 创建文本处理管道
            self.model = Pipeline([
                ('tfidf', TfidfVectorizer(
                    tokenizer=self._tokenize,
                    lowercase=True,
                    max_features=5000,
                    ngram_range=(1, 2),
                    stop_words=self._get_stop_words()
                )),
                ('classifier', MultinomialNB(alpha=0.1))
            ])
            
            # 使用预定义的训练数据进行训练
            self._train_with_sample_data()
            
            audit_logger.info("文本分类模型初始化成功")
            
        except Exception as e:
            audit_logger.error(f"文本分类模型初始化失败: {e}")
    
    def _tokenize(self, text: str) -> List[str]:
        """中文分词"""
        # 清理文本
        text = re.sub(r'[^\w\s\u4e00-\u9fff]', ' ', text)
        # 使用jieba分词
        tokens = jieba.lcut(text)
        # 过滤短词和数字
        tokens = [token for token in tokens if len(token) > 1 and not token.isdigit()]
        return tokens
    
    def _get_stop_words(self) -> List[str]:
        """获取停用词列表"""
        return [
            '的', '了', '在', '是', '我', '有', '和', '就', '不', '人',
            '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去',
            '你', '会', '着', '没有', '看', '好', '自己', '这', '那', '什么'
        ]
    
    def _train_with_sample_data(self):
        """使用样本数据训练模型"""
        # 样本训练数据
        sample_texts = [
            # 正常内容
            "今天天气真好，适合出去走走",
            "期末考试快到了，大家一起加油",
            "图书馆的学习氛围很好",
            "食堂的饭菜还不错",
            "这门课程很有意思",
            
            # 垃圾信息
            "免费领取现金红包，点击链接",
            "刷单兼职，日赚300元",
            "加微信群，免费送礼品",
            "投资理财，月收益20%",
            "办证刻章，快速办理",
            
            # 广告推广
            "代购正品化妆品，价格优惠",
            "微商招代理，月入过万",
            "淘宝店铺推广，刷好评",
            "培训班招生，包过包会",
            "二手物品出售，价格便宜",
            
            # 不当内容
            "学校管理太严格了，真讨厌",
            "老师讲课太无聊，想睡觉",
            "考试太难了，想要答案",
            "室友太吵了，很烦人",
            "食堂饭菜太难吃了"
        ]
        
        sample_labels = [
            'normal', 'normal', 'normal', 'normal', 'normal',
            'spam', 'spam', 'spam', 'spam', 'spam',
            'advertisement', 'advertisement', 'advertisement', 'advertisement', 'advertisement',
            'inappropriate', 'inappropriate', 'inappropriate', 'inappropriate', 'inappropriate'
        ]
        
        # 训练模型
        self.model.fit(sample_texts, sample_labels)
        audit_logger.info("文本分类模型训练完成")
    
    def classify(self, text: str) -> TextClassificationResult:
        """对文本进行分类"""
        try:
            if not text or not self.model:
                return TextClassificationResult(
                    category='normal',
                    confidence=0.5,
                    is_normal=True
                )
            
            # 预测分类
            prediction = self.model.predict([text])[0]
            
            # 获取预测概率
            probabilities = self.model.predict_proba([text])[0]
            classes = self.model.classes_
            
            # 找到预测类别的概率
            pred_index = np.where(classes == prediction)[0][0]
            confidence = probabilities[pred_index]
            
            # 判断是否为正常内容
            is_normal = prediction == 'normal'
            
            # 使用规则增强判断
            enhanced_result = self._enhance_with_rules(text, prediction, confidence)
            
            result = TextClassificationResult(
                category=enhanced_result['category'],
                confidence=enhanced_result['confidence'],
                is_normal=enhanced_result['is_normal']
            )
            
            if not is_normal:
                audit_logger.info(f"文本分类结果: {prediction}, 置信度: {confidence:.2f}")
            
            return result
            
        except Exception as e:
            audit_logger.error(f"文本分类失败: {e}")
            return TextClassificationResult(
                category='normal',
                confidence=0.0,
                is_normal=True
            )
    
    def _enhance_with_rules(self, text: str, prediction: str, confidence: float) -> Dict:
        """使用规则增强分类结果"""
        # 广告关键词
        ad_keywords = ['代购', '微商', '兼职', '刷单', '加微信', '免费', '赚钱', '投资', '理财']
        # 垃圾信息关键词
        spam_keywords = ['点击链接', '立即下载', '免费领取', '限时优惠', '办证', '刻章']
        # 不当内容关键词
        inappropriate_keywords = ['讨厌', '烦人', '垃圾', '傻逼', '白痴']
        
        text_lower = text.lower()
        
        # 检查广告关键词
        ad_count = sum(1 for keyword in ad_keywords if keyword in text)
        if ad_count >= 2:
            return {
                'category': 'advertisement',
                'confidence': min(confidence + 0.2, 1.0),
                'is_normal': False
            }
        
        # 检查垃圾信息关键词
        spam_count = sum(1 for keyword in spam_keywords if keyword in text)
        if spam_count >= 1:
            return {
                'category': 'spam',
                'confidence': min(confidence + 0.3, 1.0),
                'is_normal': False
            }
        
        # 检查不当内容关键词
        inappropriate_count = sum(1 for keyword in inappropriate_keywords if keyword in text)
        if inappropriate_count >= 1:
            return {
                'category': 'inappropriate',
                'confidence': min(confidence + 0.2, 1.0),
                'is_normal': False
            }

        return {
            'category': prediction,
            'confidence': confidence,
            'is_normal': prediction == 'normal'
        }

    def get_category_name(self, category: str) -> str:
        """获取分类名称"""
        return self.categories.get(category, '未知分类')


# 全局文本分类器实例
text_classifier = TextClassifier()
